
#include<src/pcps.h>
#include<src/cholesky.h>
#include<src/random.h>
#include<src/blas_lapack.h>

template<class T> class CholeskySourceTester : public pcps::CholeskySource<T> {

  private:

    int _n;
    double _thresh;
    std::vector<T> _A;

  public:

    std::vector<T> & A() { return _A; }

    int chol_size() const { return _n; }

    double chol_thresh() const { return _thresh; }

    double chol_diag(const int i) const { return pcps::real(_A[i*_n+i]); }

    void chol_get_col(const int j, std::vector<T> & col) const {
      assert( col.size() == _n );
      for (int i = 0; i < _n; i++)
        col[i] = _A[i*_n+j];
    }

    CholeskySourceTester(const int n, const double thresh)
      : _n(n),
        _thresh(thresh),
        _A(n*n, pcps::zero<T>())
    {

      for (int i = 0; i < n; i++)
      for (int j = 0; j < n; j++) {
        _A[i*n+j] = 0.01 * pcps::random_number<T>() - 0.005*(pcps::unity<T>() + pcps::imaginary_unity<T>());
        _A[j*n+i] = pcps::conj(_A[i*n+j]);
      }
      for (int i = 0; i < n; i++)
        _A[i*n+i] = pcps::unity<T>() * ( 10.0 + 5.0 * pcps::random_number<double>() );
      _A[1*n+1] = pcps::unity<T>() * ( 0.5 + 0.5 * pcps::random_number<double>() );
      _A[3*n+3] = pcps::unity<T>() * ( 0.08 + 0.05 * pcps::random_number<double>() );

    }

};

template<class T> class CholeskyResultTester : public pcps::CholeskyTarget<T> {

  private:

    int _n;
    std::vector<T> _Q;

  public:

    std::vector<T> & Q() { return _Q; }

    CholeskyResultTester(const int n) : _n(n), _Q(n*n, pcps::zero<T>()) {}

    void chol_prepare(const int m) {}

    void chol_put_col(const int j, const std::vector<T> & col) {
      assert( col.size() == _n );
      for (int i = 0; i < _n; i++)
        _Q[i*_n+j] = col[i];
    }

};

template<class T> void pcps::test_cholesky(const int n) {

  CholeskySourceTester<T> s(n, 0.01);
  CholeskyResultTester<T> t(n);

  pcps::cholesky_decomp(s, t);

  std::cout << boost::format("A:") << std::endl;
  for (int i = 0; i < n; i++) {
    for (int j = 0; j < n; j++)
      std::cout << pcps::format_number("  %10.7f", s.A()[i*n+j]);
    std::cout << std::endl;
  }
  std::cout << std::endl;

  std::cout << boost::format("Q:") << std::endl;
  for (int i = 0; i < n; i++) {
    for (int j = 0; j < n; j++)
      std::cout << pcps::format_number("  %10.7f", t.Q()[i*n+j]);
    std::cout << std::endl;
  }
  std::cout << std::endl;

  for (int i = 0; i < n; i++)
  for (int j = 0; j < n; j++)
  for (int k = 0; k < n; k++)
    s.A()[i*n+j] -= t.Q()[i*n+k] * pcps::conj(t.Q()[j*n+k]);

  double error = 0.0;
  for (int i = 0; i < s.A().size(); i++)
    error += std::abs(s.A()[i]);

  std::cout << boost::format("cholesky decomposition error = %.2e") % error << std::endl;
  std::cout << std::endl;

}

template void pcps::test_cholesky<double>(const int);
template void pcps::test_cholesky<std::complex<double> >(const int);
